#pragma once

#include <unifex/config.hpp>
#include <unifex/async_trace.hpp>
#include <unifex/bind_back.hpp>
#include <unifex/receiver_concepts.hpp>
#include <unifex/sender_concepts.hpp>
#include <unifex/std_concepts.hpp>
#include <unifex/tag_invoke.hpp>
#include <unifex/type_list.hpp>
#include <unifex/type_traits.hpp>

#include <functional>

#include <unifex/detail/prologue.hpp>

namespace unifex {
namespace _upon_error {
namespace detail {
template <typename Result, typename = void>
struct result_overload {
  using type = type_list<Result>;
};
template <>
struct result_overload<void> {
  using type = type_list<>;
};

template <typename Result>
using result_overload_t = typename result_overload<Result>::type;

}  // namespace detail

template <typename Receiver, typename Func>
struct _receiver {
  struct type;
};
template <typename Receiver, typename Func>
using receiver_t = typename _receiver<Receiver, Func>::type;

template <typename Receiver, typename Func>
struct _receiver<Receiver, Func>::type {
  UNIFEX_NO_UNIQUE_ADDRESS Func func_;
  UNIFEX_NO_UNIQUE_ADDRESS Receiver receiver_;

  template(typename... Values)  //
      (requires receiver_of<
          Receiver,
          Values...>)  //
      void set_value(Values&&... values) && {
    unifex::set_value((Receiver &&)(receiver_), (Values &&)(values)...);
  }

  template(typename Error)                  //
      (requires receiver<Receiver, Error>)  //
      void set_error(Error&& error) && noexcept {
    using result_t = std::invoke_result_t<Func, Error>;
    if constexpr (std::is_void_v<result_t>) {
      if constexpr (noexcept(std::invoke((Func &&) func_, (Error &&) error))) {
        std::invoke((Func &&) func_, (Error &&) error);
        unifex::set_value((Receiver &&) receiver_);
      } else {
        UNIFEX_TRY {
          std::invoke((Func &&) func_, (Error &&) error);
          unifex::set_value((Receiver &&) receiver_);
        }
        UNIFEX_CATCH(...) {
          unifex::set_error((Receiver &&) receiver_, std::current_exception());
        }
      }
    } else {
      if constexpr (noexcept(std::invoke((Func &&) func_, (Error &&) error))) {
        unifex::set_value(
            (Receiver &&) receiver_,
            std::invoke((Func &&) func_, (Error &&) error));
      } else {
        UNIFEX_TRY {
          unifex::set_value(
              (Receiver &&) receiver_,
              std::invoke((Func &&) func_, (Error &&) error));
        }
        UNIFEX_CATCH(...) {
          unifex::set_error((Receiver &&) receiver_, std::current_exception());
        }
      }
    }
  }

  void set_done() && noexcept { unifex::set_done((Receiver &&)(receiver_)); }

  template(typename CPO)                       //
      (requires is_receiver_query_cpo_v<CPO>)  //
      friend auto tag_invoke(CPO cpo, const type& r) noexcept(
          is_nothrow_callable_v<CPO, const Receiver&>)
          -> callable_result_t<CPO, const Receiver&> {
    return (CPO &&)(cpo)(std::as_const(r.receiver_));
  }

#if UNIFEX_ENABLE_CONTINUATION_VISITATIONS
  template <typename Visit>
  friend void
  tag_invoke(tag_t<visit_continuations>, const type& self, Visit&& visit) {
    std::invoke(visit, self.receiver_);
  }
#endif
};

template <typename Predecessor, typename Func>
struct _sender {
  struct type;
};
template <typename Predecessor, typename Func>
using sender =
    typename _sender<remove_cvref_t<Predecessor>, std::decay_t<Func>>::type;

template <typename Predecessor, typename Func>
struct _sender<Predecessor, Func>::type {
  UNIFEX_NO_UNIQUE_ADDRESS Predecessor pred_;
  UNIFEX_NO_UNIQUE_ADDRESS Func func_;

private:
  /*
   * This helper returns type_list<type_list<Result>> if invoking func returns
   * Result else if func returns void then helper returns type_list<type_list<>>
   */
  template <typename... Error>
  using invoked_result_t = type_list<
      detail::result_overload_t<std::invoke_result_t<Func, Error>>...>;

public:
  template <
      template <typename...>
      class Variant,
      template <typename...>
      class Tuple>
  using value_types = type_list_nested_apply_t<
      concat_type_lists_unique_t<
          sender_value_types_t<Predecessor, type_list, type_list>,
          sender_error_types_t<Predecessor, invoked_result_t>>,
      Variant,
      Tuple>;

  template <template <typename...> class Variant>
  using error_types = type_list<std::exception_ptr>::apply<Variant>;

  static constexpr bool sends_done = sender_traits<Predecessor>::sends_done;

  static constexpr blocking_kind blocking =
      sender_traits<Predecessor>::blocking;

  static constexpr bool is_always_scheduler_affine =
      sender_traits<Predecessor>::is_always_scheduler_affine;

  template <typename Receiver>
  using receiver_t = receiver_t<Receiver, Func>;

  friend constexpr blocking_kind
  tag_invoke(tag_t<blocking>, const type& sender) noexcept {
    return unifex::blocking(sender.pred_);
  }

  template(typename Sender, typename Receiver)  //
      (requires same_as<remove_cvref_t<Sender>, type> AND receiver<Receiver> AND
           sender_to<
               member_t<Sender, Predecessor>,
               receiver_t<remove_cvref_t<Receiver>>>)  //
      friend auto tag_invoke(
          tag_t<unifex::connect>,
          Sender&& s,
          Receiver&&
              r) noexcept(std::
                              is_nothrow_constructible_v<
                                  remove_cvref_t<Receiver>,
                                  Receiver>&&
                                  std::is_nothrow_constructible_v<
                                      Func,
                                      member_t<Sender, Func>>&&
                                      is_nothrow_connectable_v<
                                          member_t<Sender, Predecessor>,
                                          receiver_t<remove_cvref_t<Receiver>>>)
          -> connect_result_t<
              member_t<Sender, Predecessor>,
              receiver_t<remove_cvref_t<Receiver>>> {
    return unifex::connect(
        static_cast<Sender&&>(s).pred_,
        receiver_t<remove_cvref_t<Receiver>>{
            static_cast<Sender&&>(s).func_, static_cast<Receiver&&>(r)});
  }
};

namespace _cpo {
struct _fn {
private:
  template <typename Sender, typename Func>
  using _result_t = typename conditional_t<
      tag_invocable<_fn, Sender, Func>,
      meta_tag_invoke_result<_fn>,
      meta_quote2<_upon_error::sender>>::template apply<Sender, Func>;

public:
  template(typename Sender, typename Func)         //
      (requires tag_invocable<_fn, Sender, Func>)  //
      auto
      operator()(Sender&& predecessor, Func&& func) const
      noexcept(is_nothrow_tag_invocable_v<_fn, Sender, Func>)
          -> _result_t<Sender, Func> {
    return unifex::tag_invoke(_fn{}, (Sender &&)(predecessor), (Func &&)(func));
  }

  template(typename Sender, typename Func)           //
      (requires(!tag_invocable<_fn, Sender, Func>))  //
      auto
      operator()(Sender&& predecessor, Func&& func) const
      noexcept(std::is_nothrow_constructible_v<
               _upon_error::sender<Sender, Func>,
               Sender,
               Func>) -> _result_t<Sender, Func> {
    return _upon_error::sender<Sender, Func>{
        (Sender &&)(predecessor), (Func &&)(func)};
  }

  template <typename Func>
  constexpr auto operator()(Func&& func) const
      noexcept(is_nothrow_callable_v<tag_t<bind_back>, _fn, Func>)
          -> bind_back_result_t<_fn, Func> {
    return bind_back(*this, (Func &&)(func));
  }
};
}  // namespace _cpo
}  // namespace _upon_error
inline constexpr _upon_error::_cpo::_fn upon_error{};
}  // namespace unifex

#include <unifex/detail/epilogue.hpp>
